package ospf

import (
	"ConvertFlow/message"
	"ConvertFlow/utils"
	"encoding/binary"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pcapgo"
	"github.com/sirupsen/logrus"
	"log"
	"net"
	"os"
	"time"
)

func Chunks(s string, chunkSize int) []string {
	if len(s) == 0 {
		return nil
	}
	if chunkSize >= len(s) {
		return []string{s}
	}
	var chunks = make([]string, 0, (len(s)-1)/chunkSize+1)
	currentLen := 0
	currentStart := 0
	for i := range s {
		if currentLen == chunkSize {
			chunks = append(chunks, s[currentStart:i])
			currentLen = 0
			currentStart = i
		}
		currentLen++
	}
	chunks = append(chunks, s[currentStart:])
	return chunks
}

// 注册自定义层类型
var (
	LayerTypeOSPFv2 = gopacket.RegisterLayerType(
		2001,
		gopacket.LayerTypeMetadata{
			Name:    "OSPFv2",
			Decoder: gopacket.DecodeFunc(decodeOSPFv2),
		},
	)
	LayerTypeOSPFv2Hello = gopacket.RegisterLayerType(
		2002,
		gopacket.LayerTypeMetadata{
			Name:    "OSPFv2Hello",
			Decoder: gopacket.DecodeFunc(decodeOSPFv2Hello),
		},
	)
)

// OSPFv2 包结构体
type OSPFv2 struct {
	layers.BaseLayer
	Version        uint8
	Type           uint8
	PacketLength   uint16
	RouterID       uint32
	AreaID         uint32
	Checksum       uint16
	AuType         uint16
	Authentication uint64
	Contents       []byte // 存储除头部外的数据
}

// OSPFv2Hello 包结构体
type OSPFv2Hello struct {
	layers.BaseLayer
	NetworkMask   uint32
	HelloInterval uint16
	Options       uint8
	Priority      uint8
	DeadInterval  uint32
	DR            uint32
	BDR           uint32
	Neighbors     []uint32
}

// OSPF客户端结构体
type OSPFClient struct {
	handle    *pcap.Handle
	localIP   net.IP
	remoteIP  net.IP
	routerID  uint32
	Interface string
	SrcMAC    net.HardwareAddr
	DstMAC    net.HardwareAddr
}

// LayerType 实现
func (o *OSPFv2) LayerType() gopacket.LayerType      { return LayerTypeOSPFv2 }
func (h *OSPFv2Hello) LayerType() gopacket.LayerType { return LayerTypeOSPFv2Hello }

// 解码函数实现
func decodeOSPFv2(data []byte, p gopacket.PacketBuilder) error {
	if len(data) < OSPFv2HeaderLength {
		return fmt.Errorf("OSPF packet too short")
	}

	ospf := &OSPFv2{
		Version:        data[0],
		Type:           data[1],
		PacketLength:   binary.BigEndian.Uint16(data[2:4]),
		RouterID:       binary.BigEndian.Uint32(data[4:8]),
		AreaID:         binary.BigEndian.Uint32(data[8:12]),
		Checksum:       binary.BigEndian.Uint16(data[12:14]),
		AuType:         binary.BigEndian.Uint16(data[14:16]),
		Authentication: binary.BigEndian.Uint64(data[16:24]),
		Contents:       data[24:],
	}

	// 验证包长度
	if uint16(len(data)) != ospf.PacketLength {
		return fmt.Errorf("invalid OSPF packet length")
	}

	// 验证校验和
	if !validateChecksum(data) {
		return fmt.Errorf("invalid OSPF checksum")
	}

	ospf.BaseLayer = layers.BaseLayer{Contents: data[:24], Payload: data[24:]}
	p.AddLayer(ospf)

	// 根据类型解码后续内容
	switch ospf.Type {
	case OSPFHello:
		return decodeOSPFv2Hello(ospf.Contents, p)
	}

	return nil
}

func decodeOSPFv2Hello(data []byte, p gopacket.PacketBuilder) error {
	if len(data) < HelloPacketMinLength {
		return fmt.Errorf("OSPF Hello packet too short")
	}

	hello := &OSPFv2Hello{
		NetworkMask:   binary.BigEndian.Uint32(data[0:4]),
		HelloInterval: binary.BigEndian.Uint16(data[4:6]),
		Options:       data[6],
		Priority:      data[7],
		DeadInterval:  binary.BigEndian.Uint32(data[8:12]),
		DR:            binary.BigEndian.Uint32(data[12:16]),
		BDR:           binary.BigEndian.Uint32(data[16:20]),
	}

	// 解析邻居列表
	remainingBytes := len(data) - HelloPacketMinLength
	if remainingBytes > 0 && remainingBytes%4 == 0 {
		neighborCount := remainingBytes / 4
		hello.Neighbors = make([]uint32, neighborCount)
		for i := 0; i < neighborCount; i++ {
			start := HelloPacketMinLength + (i * 4)
			hello.Neighbors[i] = binary.BigEndian.Uint32(data[start : start+4])
		}
	}

	hello.BaseLayer = layers.BaseLayer{Contents: data}
	p.AddLayer(hello)

	return nil
}

// SerializeTo 实现
func (o *OSPFv2) SerializeTo(b gopacket.SerializeBuffer, opts gopacket.SerializeOptions) error {
	bytes, err := b.PrependBytes(OSPFv2HeaderLength)
	if err != nil {
		return err
	}

	// 填充头部
	bytes[0] = o.Version
	bytes[1] = o.Type
	if opts.FixLengths {
		o.PacketLength = uint16(OSPFv2HeaderLength + len(o.Contents))
	}
	binary.BigEndian.PutUint16(bytes[2:4], o.PacketLength)
	binary.BigEndian.PutUint32(bytes[4:8], o.RouterID)
	binary.BigEndian.PutUint32(bytes[8:12], o.AreaID)
	// 校验和位置先填0
	binary.BigEndian.PutUint16(bytes[12:14], 0)
	binary.BigEndian.PutUint16(bytes[14:16], o.AuType)
	binary.BigEndian.PutUint64(bytes[16:24], o.Authentication)

	// 添加内容
	if len(o.Contents) > 0 {
		payload, err := b.AppendBytes(len(o.Contents))
		if err != nil {
			return err
		}
		copy(payload, o.Contents)
	}

	// 计算校验和
	if opts.ComputeChecksums {
		checksum := calculateFletcherChecksum(b.Bytes())
		binary.BigEndian.PutUint16(bytes[12:14], checksum)
	}

	return nil
}

func (h *OSPFv2Hello) SerializeTo(b gopacket.SerializeBuffer, opts gopacket.SerializeOptions) error {
	totalLength := HelloPacketMinLength + (len(h.Neighbors) * 4)
	// 确保4字节对齐
	if padding := totalLength % 4; padding != 0 {
		totalLength += 4 - padding
	}

	bytes, err := b.PrependBytes(totalLength)
	if err != nil {
		return err
	}

	// 清零
	for i := range bytes {
		bytes[i] = 0
	}

	// 填充固定字段
	binary.BigEndian.PutUint32(bytes[0:4], h.NetworkMask)
	binary.BigEndian.PutUint16(bytes[4:6], h.HelloInterval)
	bytes[6] = h.Options
	bytes[7] = h.Priority
	binary.BigEndian.PutUint32(bytes[8:12], h.DeadInterval)
	binary.BigEndian.PutUint32(bytes[12:16], h.DR)
	binary.BigEndian.PutUint32(bytes[16:20], h.BDR)

	// 填充邻居列表
	for i, neighbor := range h.Neighbors {
		binary.BigEndian.PutUint32(bytes[20+i*4:24+i*4], neighbor)
	}

	return nil
}

// Fletcher校验和计算
func calculateFletcherChecksum(data []byte) uint16 {
	// RFC 1008中定义的Fletcher校验和
	var c0, c1 uint16
	for i, b := range data {
		// 跳过校验和字段
		if i != 12 && i != 13 {
			c0 = (c0 + uint16(b)) % 255
			c1 = (c1 + c0) % 255
		}
	}
	return (c1 << 8) | c0
}

func validateChecksum(data []byte) bool {
	// 保存原始校验和
	originalChecksum := binary.BigEndian.Uint16(data[12:14])

	// 计算校验和时将校验和字段置0
	binary.BigEndian.PutUint16(data[12:14], 0)

	// 计算校验和
	calculatedChecksum := calculateFletcherChecksum(data)

	// 恢复原始校验和
	binary.BigEndian.PutUint16(data[12:14], originalChecksum)

	return originalChecksum == calculatedChecksum
}

// 创建新的OSPF客户端
func NewOSPFClient(iface string, localIP, remoteIP string, routerID uint32, srcMAC, dstMAC string) (*OSPFClient, error) {
	handle, err := pcap.OpenLive(iface, 65535, true, pcap.BlockForever)
	if err != nil {
		return nil, fmt.Errorf("failed to open interface: %v", err)
	}

	// 解析MAC地址
	srcHW, err := net.ParseMAC(srcMAC)
	if err != nil {
		handle.Close()
		return nil, fmt.Errorf("invalid source MAC address: %v", err)
	}

	dstHW, err := net.ParseMAC(dstMAC)
	if err != nil {
		handle.Close()
		return nil, fmt.Errorf("invalid destination MAC address: %v", err)
	}

	return &OSPFClient{
		handle:    handle,
		localIP:   net.ParseIP(localIP),
		remoteIP:  net.ParseIP(remoteIP),
		routerID:  routerID,
		Interface: iface,
		SrcMAC:    srcHW,
		DstMAC:    dstHW,
	}, nil
}

// 发送OSPF Hello包
func (c *OSPFClient) SendHello(data uint32) error {
	// 构建以太网层
	eth := &layers.Ethernet{
		SrcMAC:       c.SrcMAC,
		DstMAC:       c.DstMAC,
		EthernetType: layers.EthernetTypeIPv4,
	}

	// 构建IP层 - 添加完整字段
	ip := &layers.IPv4{
		Version:    4,
		IHL:        5,
		TOS:        0xc0, // OSPF使用的Type of Service
		Length:     0,    // 会被SerializeLayers自动计算
		Id:         0,    // 可以使用随机ID
		Flags:      0,
		FragOffset: 0,
		TTL:        1,
		Protocol:   89, // OSPF
		SrcIP:      c.localIP,
		DstIP:      c.remoteIP,
		Options:    []layers.IPv4Option{}, // 空的IP选项切片
		Padding:    []byte{},              // 空的填充字节切片
		Checksum:   0,                     // 会被SerializeLayers自动计算
	}

	// 构建OSPF Hello包
	hello := &OSPFv2Hello{
		NetworkMask:   binary.BigEndian.Uint32(net.IPv4Mask(255, 255, 255, 0)),
		HelloInterval: 10,
		Options:       OSPFOptE, // 启用外部路由能力
		Priority:      1,
		DeadInterval:  data, // 隐藏数据
		DR:            0,
		BDR:           0,
		Neighbors:     []uint32{},
	}

	// 序列化Hello包
	helloBuffer := gopacket.NewSerializeBuffer()
	err := hello.SerializeTo(helloBuffer, gopacket.SerializeOptions{
		FixLengths: true,
	})
	if err != nil {
		return fmt.Errorf("failed to serialize Hello packet: %v", err)
	}

	// 构建OSPF头
	ospf := &OSPFv2{
		Version:  2,
		Type:     OSPFHello,
		RouterID: c.routerID,
		AreaID:   0,
		AuType:   OSPFNoAuth,
		Contents: helloBuffer.Bytes(),
	}

	// 序列化完整包
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}

	err = gopacket.SerializeLayers(buf, opts,
		eth,  // 添加以太网层
		ip,   //添加ip层
		ospf, //添加ospf层
	)
	if err != nil {
		return fmt.Errorf("failed to serialize packet: %v", err)
	}

	// 发送包
	packetData := buf.Bytes()

	// 添加调试信息
	log.Printf("Packet length: %d", len(packetData))
	log.Printf("IP header: %+v", ip)

	// 使用gopacket解析生成的数据包进行验证
	packet := gopacket.NewPacket(packetData, layers.LayerTypeEthernet, gopacket.Default)
	if ipLayer := packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
		ipPacket, _ := ipLayer.(*layers.IPv4)
		log.Printf("Generated IP packet: %+v", ipPacket)
	}

	err = c.handle.WritePacketData(packetData)
	if err != nil {
		return fmt.Errorf("failed to send packet: %v", err)
	}

	// 保存为pcap文件
	filename := fmt.Sprintf("ospf_packet_%d.pcap", time.Now().Unix())
	f, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("failed to create pcap file: %v", err)
	}
	defer f.Close()

	// 创建pcap写入器，使用LinkTypeEthernet
	w := pcapgo.NewWriter(f)
	err = w.WriteFileHeader(65535, layers.LinkTypeEthernet)
	if err != nil {
		return fmt.Errorf("failed to write pcap header: %v", err)
	}

	// 写入数据包
	err = w.WritePacket(gopacket.CaptureInfo{
		Timestamp:     time.Now(),
		CaptureLength: len(packetData),
		Length:        len(packetData),
	}, packetData)
	if err != nil {
		return fmt.Errorf("failed to write packet to pcap: %v", err)
	}

	log.Printf("Packet saved to %s", filename)
	return nil
}

// 关闭客户端
func (c *OSPFClient) Close() {
	if c.handle != nil {
		c.handle.Close()
	}
}

// SendNoReply 无脑发送，无需等待对端回应
func SendNoReply(Interface string, srcMac string, dstMac string, listeningReplyAddr string, remoteAddr string, chunkSize int, delay int, data string) {
	//TODO:根据选择的隐藏形式，来选择数据的大小，比如hello能隐藏的数据大小为4个字节，则chunkSize为4
	chunkSize = 2                        //用于测试，先写死
	dataSlice := Chunks(data, chunkSize) //1 character = 1byte , max size of icmp data 65507
	dataSlice = message.Marshall(dataSlice)
	time.Sleep(time.Duration(delay) * time.Second)

	//创建OSPF客户端对象
	client, err := NewOSPFClient(
		Interface,
		listeningReplyAddr,
		remoteAddr,
		binary.BigEndian.Uint32(net.ParseIP("192.168.0.1").To4()),
		srcMac,
		dstMac,
	)
	if err != nil {
		logrus.Errorf("Failed to create OSPF client: %v", err)
	}
	defer client.Close()

	// 告知对方数据包数量,比如是1个还是2个
	testData := uint32(len(dataSlice))
	err = client.SendHello(testData)
	if err != nil {
		logrus.Error("send data packet count failed:", err.Error())
	}

	//发送实际数据包,以指定的时间间隔
	for i := 0; i < len(dataSlice); i++ {
		time.Sleep(time.Duration(delay) * time.Second)
		//TODO:fix bug当数据过大时，可能导致索引index的范围为(0-9)(10-19)(20-99等)最多支持拆分为100份进行传输
		d, err := utils.StringToUint32(dataSlice[i])
		if err != nil {
			logrus.Error("utils.StringToUint32 function failed ", err.Error())
			continue
		}
		//在发送之前会统一处理好大小端的问题,uint32类型的值无需转换大端和小端了
		err = client.SendHello(uint32(d))
		if err != nil {
			log.Println("client.SendHello function failed ", err.Error())
			continue
		}
	}
}
